# 写入lib.rs文件
# test.rs文件中写入use
# 测试文件中重新封装test函数
import json
import os
import logging
from split_function import split_rust_functions_in_file
from extract_component import extract_function_names, extract_static_declaration


# 文件路径
test_func="""
#[test]
fn {func_name}() {{
    unsafe {{
        {content}
    }}
}}
"""
test_ignore_func="""
#[test]
#[ignore]
fn {func_name}() {{
    unsafe {{
        {content}
    }}
}}
"""

# test_func_2="""

#     #[test]
#     fn {func_name}() {{
#         unsafe {{
#             {content}
#         }}
#     }}
# """

# test_ignore_func_2="""
#     #[test]
#     #[ignore]
#     fn {func_name}() {{
#         unsafe {{
#             {content}
#         }}
#     }}
# """

# 将文件中的指定代码注释掉（每行前加//) or 直接移除
# 返回F/T
def comment_code(file_path, target_code):
    with open(file_path,'r') as file:
        content = file.read()
    if target_code not in content:
        logging.info(f"未能在文件{file_path}中找到目标代码：\n{target_code}")
        return False
    commented_code = '\n'.join(f'// {line}' for line in target_code.split('\n'))
    with open(file_path,'w') as file:
        # file.write(content.replace(target_code, commented_code))
        file.write(content.replace(f"{target_code}\n", ''))
    return True



def add_stest_functions(path):
    PROJ_DIR = path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")


    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            func_names = extract_function_names(file_path)

            if func_names == None:
                logging.info(f"文件{file_path}添加安全测试函数失败")
                return

            logging.info(f"add s_test for {file_path}")
            test_func_names = [f for f in func_names if "test_" in f and "out_of_memory" not in f]
            # print(f"add s_test for {filename}")
            with open(file_path, 'a') as file:
                for f in test_func_names:
                    content = f"{f}();"
                    if "test_rb_tree_remove" in f or "test_rb_tree_to_array" in f:
                        func_code = test_ignore_func.format(func_name=f"s_{f}", content = content)
                    else:
                        func_code = test_func.format(func_name=f"s_{f}", content = content)
                    file.write("\n")
                    file.write(func_code)
        
                    

def remove_stest_functions(path):
    PROJ_DIR = path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    
    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            func_names = extract_function_names(file_path)

            if func_names == None:
                logging.info(f"文件{file_path}移除安全测试函数失败")
                return

            logging.info(f"remove s_test for {file_path}")
            test_func_names = [f for f in func_names if "test_" in f and "out_of_memory" not in f]

            with open(file_path, 'r') as file:
                file_content = file.read()
            
            for f in test_func_names:
                content = f'{f.replace("s_", "", 1)}();'
                if "test_rb_tree_remove" in f or "test_rb_tree_to_array" in f:
                    func_code = test_ignore_func.format(func_name=f"{f}", content = content)
                else:
                    func_code = test_func.format(func_name=f"{f}", content = content)
                
                # if "test_list" in file_path:
                #     print(func_code)
                file_content = file_content.replace("\n"+func_code, "")
                file_content = file_content.replace(func_code, "")

            with open(file_path, 'w') as file:
                file.write(file_content)



def add_test_annotation(path):
    PROJ_DIR = path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            # print(file_path)
            functions = split_rust_functions_in_file(file_path)
            if functions == None:
                logging.info(f"文件{file_path}添加测试注解失败")
                return
            with open(file_path, 'r') as file:
                file_content = file.read()
            for func in functions:
                if "test_" in func['name'] and "s_test" not in func['name']:
                    func_code = func['code']
                    if "test_rb_tree_remove" in func['name'] or "test_rb_tree_to_array" in func['name']:
                        annotated_func_code = "#[test]\n#[ignore]\n" + func_code
                    else:
                        annotated_func_code = "#[test]\n" + func_code
                    file_content = file_content.replace(func_code, annotated_func_code)
            
            with open(file_path, 'w') as file:
                file.write(file_content)
        # break

def remove_test_annotations(path):
    remove_stest_functions(path)
    PROJ_DIR = path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            functions = split_rust_functions_in_file(file_path)
            if functions == None:
                return
            
            with open(file_path, 'r') as file:
                file_content = file.read()
            file_content = file_content.replace("#[test]\n", "")
            file_content = file_content.replace("#[test]", "")
            with open(file_path, 'w') as file:
                file.write(file_content)
        

def remove_out_of_memory_test(path):
    PROJ_DIR = path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            functions = split_rust_functions_in_file(file_path)
            if functions == None:
                logging.info(f"文件{file_path}移除test out of memory失败")
                return
            functions = [func for func in functions if "out_of_memory" in func['name']]
            if functions == []:
                continue
            
            logging.info(f"remove test out of memory for {file_path}")

            with open(file_path, 'r') as file:
                file_content = file.read()
            for func in functions:
                func_code = func['code']
                file_content = file_content.replace("#[no_mangle]\n"+func_code, '')
                file_content = file_content.replace(func_code, '')
            
            with open(file_path, 'w') as file:
                file.write(file_content)

def remove_main_and_tests(path):
    PROJ_DIR = path
    TEST_DIR = os.path.join(PROJ_DIR, "tests")
    for root, dirs, files in os.walk(TEST_DIR):
        for file in files:
            file_path = os.path.join(root, file)
            # print(file_path)
            functions = split_rust_functions_in_file(file_path)
            for func in functions:
                if "main" in func['name']:
                    func_code = func['code']
                    comment_code(file_path, func_code)

            res = extract_static_declaration(file_path)
            static_codes = res if res else []
            for static in static_codes:
                # print(static)
                if "UnitTestFunction" in static:
                    comment_code(file_path, static)


if __name__ == "__main__":
    # add_stest_functions("Output/primary")
    remove_stest_functions("Output/primary")
    # add_test_annotation("Output/primary")
    # remove_test_annotations("Output/primary")
    # remove_out_of_memory_test("Output/primary")



